module Tools where

import DataTypes
import StationManagement
import TrainManagement
import GenericFunctions

import Data.Maybe
import Data.List
import Data.Ord
import Data.Time.Format
import Data.Graph
import Control.Monad
import System.Locale

checkArrivals _ _ [] _ i = -1
checkArrivals sid1 sid2 ((DeparturesItem _ _ id _):depatures) ok i
    | ok == True && sid2 == id = i
    | ok == True && sid2 /= id = checkDepatures sid1 sid2 depatures True (i+1)
    | ok == False && sid1 == id = checkDepatures sid1 sid2 depatures True (i+1)
    | otherwise = checkDepatures sid1 sid2 depatures False (i+1)

checkDepatures _ _ [] _ i = -1
checkDepatures sid1 sid2 ((DeparturesItem _ _ id _):depatures) ok i
    | ok == True && sid2 == id = i
    | ok == True && sid2 /= id = checkDepatures sid1 sid2 depatures True i
    | ok == False && sid1 == id = checkDepatures sid1 sid2 depatures True i
    | otherwise = checkDepatures sid1 sid2 depatures False (i+1)

getDeparturesString:: Int -> Int -> [DeparturesItem] -> IO [Char]
getDeparturesString sid1 sid2 departures = do let index = checkDepatures sid1 sid2 departures False 0
                                              if (index == -1) then
                                                 return ""
                                              else do let departure = departures !! index
                                                          time = DataTypes.departure departure in
                                                          return ((timeStr time) ++ ", ")

getArrivalsString sid1 sid2 departures = do let index = checkArrivals sid1 sid2 departures False 0
                                            if (index == -1) then
                                                return ""
                                            else do let departure = departures !! index
                                                        time = DataTypes.arrival departure in
                                                        return ((timeStr time) ++ ", ")

getStationsConnection:: Int -> Int -> Train -> IO [Char]                                             
getStationsConnection s1 s2 (Train _ _ _ _ departures) = do departuresThere <- mapM (\x -> getDeparturesString s1 s2 (fst x)) departures  
                                                            departuresFrom <- mapM (\x -> getDeparturesString s1 s2 (snd x)) departures
                                                            let dep = departuresThere ++ departuresFrom
                                                            arrivalsThere <- mapM (\x -> getArrivalsString s1 s2 (fst x)) departures
                                                            arrivalsFrom <- mapM (\x -> getArrivalsString s1 s2 (snd x)) departures
                                                            let arr = arrivalsThere ++ arrivalsFrom
                                                            return ((concat dep) ++ "Arrivals: " ++  (concat arr))

-- Function used to return a readable connection
itemString s1 s2 ts = ("train \"" ++ (name $ head ts) ++ "\": " ++ stationStr (fromJust s1) ++ " -> " ++ stationStr (fromJust s2))

-- Function used to print connections with valid number of train changes
printValidRoutes paths nodes = mapM_ (\c -> (printValidPath nodes c) >>= showItems "Connection") paths

-- Function used to print a single valid path (with proper number of train changes)
printValidPath nodes (sid:sids) = 
     let nodeTrains = getNodeTrains (getNodeById sid nodes) in do
        printValidPath' nodes nodeTrains sid (sid:sids)
        where
            printValidPath' nodes commonTrains startsid [sid] =
                do  station1 <- getStationById startsid
                    station2 <- getStationById sid
                    departuresString <- getStationsConnection startsid sid (head commonTrains)
                    return $ [(itemString station1 station2 commonTrains) ++ " | Departures: " ++ departuresString]
                                        
            printValidPath' nodes commonTrains startsid (sid:nextSid:restSids) =
                let nodeTrains = getNodeTrains (getNodeById sid nodes)
                    nextTrains = getNodeTrains (getNodeById nextSid nodes)
                    c1Trains = getCommonTrains sid nextSid nodeTrains nextTrains
                    c2Trains = intersect c1Trains commonTrains 
                in do
                    if (null c2Trains) then do     
                        station1 <- getStationById startsid
                        station2 <- getStationById sid
                        station3 <- getStationById nextSid
                        validPath <- printValidPath' nodes c1Trains sid (nextSid:restSids)
                        departuresString <- getStationsConnection startsid sid (head commonTrains)
                        return (((itemString station1 station2 commonTrains) ++ " | Departures: " ++ departuresString):validPath)  
                    else do
                        printValidPath' nodes c2Trains startsid (nextSid:restSids)
                        

-- Get trains that are common for two stations and have direct connection between them (important!)
getCommonTrains to from trains1 trains2 = filter (\(Train _ _ _ (route1, route2) _) -> (isInfixOf [to,from] route1) || (isInfixOf [to,from] route2)) $ intersect trains1 trains2

-- Get trains of a single node
getNodeTrains node = (\(graphNode, _, _) -> DataTypes.trains graphNode) node 

-- Get node with specified station id
getNodeById sid nodes = (filter (\(_, id, _) -> id == sid) nodes) !! 0 

-- Check whether path has appropriate amount of train changes
validPath _ [] _ _ = True

validPath nodes (sid:[]) changes trains =
    let nodeTrains = getNodeTrains (getNodeById sid nodes) in
        if (changes == -1) then
            False
        else let trainIndex = elemIndex True (map (\x -> x `elem` trains) nodeTrains) in
                 if (isJust trainIndex) then
                    validPath nodes [] changes trains
                 else validPath nodes [] (changes - 1) ((nodeTrains !! 0):trains)

validPath nodes (sid:(sidNext:path)) changes trains =
    let nodeTrains = getNodeTrains (getNodeById sid nodes) 
        nextNodeTrains = getNodeTrains (getNodeById sidNext nodes)
        commonTrains = getCommonTrains sid sidNext nodeTrains nextNodeTrains in
        if (changes == -1) then
           False
        else let trainIndex = elemIndex True (map (\x -> x `elem` trains) commonTrains) in
                 if (isJust trainIndex) then
                    validPath nodes (sidNext:path) changes trains
                 else validPath nodes (sidNext:path) (changes - 1) ((commonTrains !! 0):trains)

-- Check all paths between stations whether they have appropriate amount of train changes 
checkTrainChanges _ [] _ = []
checkTrainChanges nodes (path:paths) changes = if (validPath nodes path changes []) then
                                                   path:(checkTrainChanges nodes paths changes)
                                               else checkTrainChanges nodes paths changes

-- Function used to get all nodes which are connected with current node
getNodeConnections (_, _, edges) = edges 

-- Function used to check whether current node is next to target station
checkAdjacentNodes paths visited (_, _, edges) target = if (target `elem` edges) then
                                                            (visited ++ [target]):paths
                                                        else paths

-- Function used to check all stations which are next to current node in order to find all paths between start and target stations
checkAllEdges _ _ [] newPaths _ =  newPaths
checkAllEdges visited nodes (edge:edges) newPaths target =
        checkAllEdges visited nodes edges (Tools.bfs nodes (visited ++ [edge]) newPaths target) target

-- Function used to find all connections between a given pair of nodes
bfs:: [(GraphNode, Int, [Int])] -> [Int] -> [[Int]] -> Int ->  [[Int]]
bfs nodes visited paths target  = 
       let current = (filter (\(_, sid, _) -> (last visited) == sid) nodes) !! 0
           newPaths = checkAdjacentNodes paths visited current target 
           connections = filter (\x -> not (x `elem` visited) && x /= target) (getNodeConnections current) in
           if (((length nodes) /= (length visited)) && (not (null connections))) then do
              checkAllEdges visited nodes connections newPaths target
           else newPaths
       
-- Function used to print a single graph node
printNode (_, sid, connections) = do print "Station id: " 
                                     print sid
                                     print "Connections"
                                     print connections

-- Function used to print graph nodes
printNodes [] = do return ()
printNodes (node:nodes) = do printNode node
                             printNodes nodes

-- Function used to add a new edge between existing vertices
addConnection ((GraphNode stationN trains), sid, edges) train nextSid  = 
                         if (isJust nextSid) then
                            ((GraphNode stationN (newTrains)), sid, (fromJust nextSid):edges)
                         else ((GraphNode stationN (newTrains)), sid, edges)
                         where newTrains | train `elem` trains = trains
                                         | otherwise = train:trains

-- Function used to modify an existing graph vertex
modifyNode nodes stationIndex train next = 
                     do  let node = nodes !! stationIndex
                         let splitNodes = splitAt stationIndex nodes
                         let newNodes = (fst splitNodes) ++ (tail (snd splitNodes))
                         let newNode = addConnection node train next
                         let finalNodes =  newNode:newNodes
                         return finalNodes

-- Function used to get index of vertex in the list of nodes
getNodeListIndex station nodes =  elemIndex station (map (\node -> DataTypes.station node) (map (\(e, _, _) -> e) nodes))

-- Function used to prepare graph vertices and edges from single train route
reflectRoute:: Train -> [Int] -> [(GraphNode, Int, [Int])] -> IO [(GraphNode, Int, [Int])]
reflectRoute _ [] nodes = do return nodes

reflectRoute train (sid:[]) nodes = 
                   do station <- getStationById sid
                      let stationIndex =  getNodeListIndex (fromJust station) nodes
                      if (isJust stationIndex) then do
                         nodes <- modifyNode nodes (fromJust stationIndex) train Nothing
                         return nodes
                      else do 
                         return (((GraphNode (fromJust station) [train]), sid, []):nodes)

reflectRoute train (sid:(nextSID:stations)) nodes = 
                do station <- getStationById sid
                   let stationIndex = getNodeListIndex (fromJust station) nodes
                   if (isJust stationIndex) then do
                      nodes <- modifyNode nodes (fromJust stationIndex) train (Just nextSID)
                      reflectRoute train ((nextSID:stations)) nodes
                   else do 
                      reflectRoute train ((nextSID:stations)) (((GraphNode (fromJust station) [train]), sid, [nextSID]):nodes)

-- Function used to prepare vertices and edges from single train
reflectTrain:: Train -> [(GraphNode, Int, [Int])] -> IO [(GraphNode, Int, [Int])]
reflectTrain train nodes = do newNodes <- reflectRoute train (fst (DataTypes.route train)) nodes
                              reflectRoute train (snd (DataTypes.route train)) newNodes

-- Function used to create vertices and edges of the graph
createNodes:: [Train] ->  [(GraphNode, Int, [Int])] -> IO [(GraphNode, Int, [Int])]
createNodes [] nodes = do return nodes
createNodes (train:trains) nodes = do newNodes <- reflectTrain train nodes
                                      createNodes trains newNodes

-- Function used to generate a train schedule
generateSchedule =  do  date <- promptDate "Enter date"
                        s <- StationManagement.promptStation
                        when (isJust s) (printSchedule (fromJust s) date)
                        
-- Function used to print a schedule
printSchedule (Station id name city) date =  
    do  trains <- TrainManagement.getTrains
        let rx = concat (map (\t -> trainDepartures t id date) trains)
        let ra = sortBy (\(t1, d1) (t2, d2) -> compare (arrival d1) (arrival d2)) rx
        mapM arrivalStr ra >>= showItems ("Arrivals to \"" ++ name ++ "\" station")
        mapM departureStr ra >>= showItems ("Departures from \"" ++ name ++ "\" station")
        pressEnter
    where
        arrivalStr (train, DeparturesItem arriv _ _ dir) = arrdepStr train arriv dir
        departureStr (train, DeparturesItem _ depart _ dir) = arrdepStr train depart dir                              
        arrdepStr train time dir = do
            st1 <- StationManagement.getStationById (head $ fst $ route train)
            st2 <- StationManagement.getStationById (last $ fst $ route train)
            st3 <- StationManagement.getStationById (last $ snd $ route train)
            let name1 = stationName $ fromJust st1
            let name2 = stationName $ fromJust st2
            let name3 = stationName $ fromJust st3        
            if (dir == Forward) then 
                return $ (timeStr time) ++ " " ++ (DataTypes.name train) ++ ", \"" ++ name1 ++ "\" -> \"" ++ name2 ++ "\""
            else
                return $ (timeStr time) ++ " " ++ (DataTypes.name train) ++ ", \"" ++ name2 ++ "\" -> \"" ++ name3 ++ "\""

-- Function used to get train departures from station specified by its id (sid) and in given date
-- returns list of tuples (Train, DepartureItem)
trainDepartures train sid date =
    let (Train tname tsince ttill (r1, r2) s) = train in
        [(train, d) | d <- (concat [s1 ++ s2 | (s1, s2) <- s]), stationId d == sid, since train <= date, till train > date]

